#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import sys
import json
import select
import shutil
import pathlib
import subprocess
import pwd
import fcntl
import errno
sys.path.append('/usr/lib64/syncupd')
from gbs_util import GbsUtil
from gbs_param import GbsConst
from gbs_param import GbsParam
from gbs_common import GbsProtocolException
from gbs_common import GbsSystem
from gbs_ctrl_session import GbsCtrlSession


class Main:

    def __init__(self):
        # these directories should be already created by syncupd main process
        assert os.path.exists(GbsConst.runDir)
        assert os.path.exists(GbsConst.varDir)
        assert os.path.exists(GbsConst.logDir)
        
        self.param = GbsParam()
        self.sessionObj = None
        self.sysObj = None
    
    def run(self):
        try:
            # get user
            user = os.environ["USER"]

            # get client ip:port
            tlist = os.environ["SSH_CLIENT"].split()
            clientIp = tlist[0]
            clientPort = tlist[1]

            # get client public key
            clientPubKey = pathlib.Path(os.environ["SSH_USER_AUTH"]).read_text()
            assert clientPubKey.startswith("publickey ")                            # the generated sshd configuration ensures that SSH_USER_AUTH exists and publickey auth must be used
            clientPubKey = clientPubKey[len("publickey "):]

            sessionType = os.environ.get("GBS_SESSION_TYPE", "")
            if sessionType == "ctrl":
                self.handleControlSession(user, clientIp, clientPort, clientPubKey)
            elif sessionType == "rsync":
                self.handleRsyncSession(clientPubKey)
            elif sessionType == "chroot":
                self.handleCommandSession(clientPubKey)
            else:
                raise Exception(f"unknown session type: {sessionType}")
        except BaseException as e:
            self.logError(f"Session failed: {e}")
            raise
        finally:
            if self.sessionObj is not None:
                self.sessionObj.dispose()
            if self.sysObj is not None:
                self.sysObj.dispose()

    def handleControlSession(self, user, clientIp, clientPort, clientPubKey):
        recvBuf = b''
        sendBuf = b''
        sessionObj = GbsCtrlSession(self.param, user, clientIp, clientPort, clientPubKey)
        self.logInfo('Control Channel: Client "%s" appeared, UUID "%s".' % (clientIp, sessionObj.sysObj.uuid))
        
        # use binary mode for stdin and stdout
        sys.stdin = sys.stdin.buffer
        sys.stdout = sys.stdout.buffer

        while True:
            # select
            inputs = [sys.stdin]
            if len(sendBuf) > 0:
                outputs = [sys.stdout]
            else:
                outputs = []
            readable, writable, exceptional = select.select(inputs, outputs, inputs, 10.0)

            # read from client, process request
            while len(readable) > 0:
                buf = sys.stdin.read(4096)
                if len(buf) == 0:
                    self.logInfo('Control Channel: Client "%s" disconnected.' % (self.formatClient()))
                    return

                recvBuf += buf
                i = recvBuf.find(b'\n')
                if i < 0:
                    # no full json object received
                    continue

                requestObj = json.loads(recvBuf[:i].decode("iso8859-1"))           # ignore "\n"
                recvBuf = recvBuf[i + 1:]                                          # jump over "\n"

                # process the received request json object
                try:
                    responseObj = sessionObj.onRequest(requestObj)                 # create response when processing request
                except GbsProtocolException as e:
                    raise Exception("Control Channel: Bad request, %s." % (str(e)))

                sendBuf += (json.dumps(responseObj) + "\n").encode("iso8859-1")

            # write to client
            if len(writable) > 0:
                i = sys.stdout.write(sendBuf)
                sys.stdout.flush()
                sendBuf = sendBuf[i:]

            # error occured
            if len(exceptional) > 0:
                raise Exception("debug2")
    
    def handleRsyncSession(self, clientPubKey):
        originalCmd = os.environ.get("SSH_ORIGINAL_COMMAND", None)
        if originalCmd is None:
            raise Exception("Rsync Channel: Original command not specified.")
        if not originalCmd.startswith("rsync --server"):
            raise Exception("Rsync Channel: Invalid original command.")

        subprocess.run(originalCmd, shell=True)
    
    def handleCommandSession(self, clientPubKey):
        originalCmd = os.environ.get("SSH_ORIGINAL_COMMAND", None)
        if originalCmd is None:
            raise Exception("Command Channel: Original command not specified.")

        # 创建GbsSystem实例
        try:
            self.sysObj = GbsSystem(self.param, clientPubKey)
            self.logInfo(f"Created GbsSystem with UUID: {self.sysObj.uuid}")
            
            # 挂载磁盘镜像
            self.sysObj.mount()
            self.logInfo(f"Mounted system image at: {self.sysObj.mntDir}")
            
            # 准备根目录
            self.sysObj.prepareRoot()
            self.logInfo("Prepared root directory for chroot")
            
            # 切换到root用户
            try:
                # 获取root用户信息
                root_pwd = pwd.getpwnam('root')
                
                # 设置环境变量
                env = os.environ.copy()
                env['HOME'] = '/root'
                env['PATH'] = '/bin:/usr/bin:/sbin:/usr/sbin'
                
                # 执行chroot命令
                self.logInfo(f"Executing chroot command: {originalCmd}")
                result = subprocess.run(
                    ['/usr/sbin/chroot', self.sysObj.mntDir, '/bin/sh', '-c', originalCmd],
                    env=env,
                    preexec_fn=lambda: os.setresuid(root_pwd.pw_uid, root_pwd.pw_gid, root_pwd.pw_gid)
                )
                
                sys.exit(result.returncode)
            except Exception as e:
                self.logError(f"Failed to execute chroot command: {e}")
                sys.exit(1)
            finally:
                # 清理根目录准备
                try:
                    if self.sysObj is not None:
                        self.sysObj.unPrepareRoot()
                except Exception as e:
                    self.logError(f"Failed to unprepare root: {e}")
                
                # 卸载磁盘镜像
                try:
                    if self.sysObj is not None:
                        self.sysObj.unmount()
                except Exception as e:
                    self.logError(f"Failed to unmount system: {e}")
        except Exception as e:
            self.logError(f"Failed to create GbsSystem: {e}")
            sys.exit(1)
    
    def logError(self, msg):
        # FIXME: send log to main process
        print(f"ERROR: {msg}", file=sys.stderr)
    
    def logInfo(self, msg):
        # FIXME: send log to main process
        print(f"INFO: {msg}", file=sys.stderr)
    
    def logDebug(self, msg):
        # FIXME: send log to main process
        print(f"DEBUG: {msg}", file=sys.stderr)
    
    def formatClient(self):
        return "unknown client"


if __name__ == "__main__":
    Main().run()
